#include "DSA.h"

#include <ctime>
#include <stdexcept>
#include <map>

#include "data_conversion.h"

#include "../DEBUG_PRINT_MACRO.h"

#ifdef DEBUG_PRINTS_LEVEL_1
#include <iostream>
#endif

// 引入 Tracy 头文件
#include <Tracy.hpp>

bool is_prime(const mpz_t& n, int reps = 25) {
    // 为 is_prime 函数添加 Tracy 区域
    ZoneScopedN("is_prime");
    return mpz_probab_prime_p(n, reps) > 0;
}

template<int key_length>
void DSA<key_length>::generatepublicparameters(DSA_publicparameter& publicParameters) {
    // 为 generatepublicparameters 函数添加 Tracy 区域
    ZoneScopedN("DSA::generatepublicparameters");

    #ifdef DEBUG_PRINTS_LEVEL_2
    std::cerr << "[DSA public parameters generate][Debug] Generating public parameters start." << std::endl;
    #endif
    // ----------------
    // 1. 选择密钥的长度
    // ----------------
    int L = key_length; // 密钥长度
    int N; // 签名的位数
    {
        ZoneScopedN("DSA::generatepublicparameters::Key Length Selection");

        std::map<int, int>N_map = { {1024,160},{2048,224},{3072,256} };
        if (N_map.find(L) != N_map.end()) {
            N = N_map[L];
        }
        else {
            throw std::runtime_error("Invalid key length");
        }
    }


    mpz_t p, q, g, k, p_minus_1, h, exp, temp;
    {
        ZoneScopedN("DSA::generatepublicparameters::GMP Variable Initialization");
        mpz_inits(p, q, g, k, p_minus_1, h, exp, temp, nullptr);
    }

    // 初始化随机数生成器
    gmp_randstate_t state;
    {
        ZoneScopedN("DSA::generatepublicparameters::Random State Initialization");
        gmp_randinit_default(state);
        gmp_randseed_ui(state, static_cast<unsigned long>(time(nullptr)));
    }


    // 最大尝试次数
    constexpr int MAX_Q_ATTEMPTS = 1000;
    constexpr int MAX_P_ATTEMPTS = 8000;
    int attempts_q = 0;
    int attempts_p = 0;
    bool success = false;

    // 生成密钥对，在不满足条件时重新生成
    while (attempts_q++ < MAX_Q_ATTEMPTS) {
        ZoneScopedN("DSA::generatepublicparameters::Loop Attempt");
        // ---------------------
        // 2. 生成 N 比特的素数 q
        // ---------------------
        {
            ZoneScopedN("DSA::generatepublicparameters::Generate Prime Q");
            do {
                mpz_urandomb(q, state, N);
                mpz_setbit(q, N - 1);   // 强制最高位为1
                mpz_nextprime(q, q);

                if (++attempts_q > MAX_Q_ATTEMPTS) {
                    #ifdef DEBUG_PRINTS_LEVEL_2
                    std::cerr << "[DSA public parameters generate][Error] Exceeded maximum number of q generation attempts." << std::endl;
                    #endif
                    throw std::runtime_error("Failed to generate proper prime q");
                }
            } while (mpz_sizeinbase(q, 2) != N);
        }
        #ifdef DEBUG_PRINTS_LEVEL_3
        gmp_printf("[DSA] Generated q: %Zd\n\n", q);
        #endif
        // ------------------------------------------
        // 3. 生成 L 比特的素数 p，使得 p-1 是 q 的倍数
        // ------------------------------------------
        {
            ZoneScopedN("DSA::generatepublicparameters::Generate Prime P");
            do {
                // 生成k值（864位，使得k*q ≈ 1024位）
                mpz_urandomb(k, state, L - N);

                // 计算p = k*q + 1
                mpz_mul(p_minus_1, q, k);
                mpz_add_ui(p, p_minus_1, 1);

                if (++attempts_p > MAX_P_ATTEMPTS) {
                    #ifdef DEBUG_PRINTS_LEVEL_2
                    std::cerr << "[DSA public parameters generate][Error] Exceeded maximum number of p generation attempts." << std::endl;
                    #endif
                    throw std::runtime_error("Failed to generate proper prime p");
                }
            } while (!is_prime(p) || mpz_sizeinbase(p, 2) != L); // is_prime 内部已插桩
        }
        #ifdef DEBUG_PRINTS_LEVEL_3
        gmp_printf("[DSA] Generated p: %Zd\n\n", p);
        #endif
        {
            ZoneScopedN("DSA::generatepublicparameters::P-1 Modulo Q Check");
            mpz_mod(temp, p_minus_1, q);
            if (mpz_cmp_ui(temp, 0) != 0) continue;
        }

        // --------------
        // 4. 选择生成元 g
        // --------------
        {
            ZoneScopedN("DSA::generatepublicparameters::Choose Generator G");
            do {
                mpz_urandomb(h, state, N);
                mpz_div(exp, p_minus_1, q);    // exp = (p-1) / q
                mpz_powm(g, h, exp, p);        // g = pow(h, (p - 1) / q, p) , h为随机数
                attempts_p++; // 每次尝试的成功率约为0.14%, 存在生成失败可能
            } while (mpz_cmp_ui(g, 1) == 0 || attempts_p >= MAX_P_ATTEMPTS); // 如果g == 1，则重新生成g
        }

        break;
    }


    // 保存公共参数
    {
        ZoneScopedN("DSA::generatepublicparameters::Store Parameters and Cleanup");
        publicParameters = DSA_publicparameter(p, q, g);
        publicParameters.is_valid = true;
        mpz_clears(p, q, g, k, p_minus_1, h, exp, temp, nullptr);
        gmp_randclear(state);
    }
}


template<int key_length>
void DSA<key_length>::generateKeyPair(DSA_privatekey& privateKey, DSA_publickey& publicKey, DSA_publicparameter& publicParameters) {
    // 为 generateKeyPair 函数添加 Tracy 区域
    ZoneScopedN("DSA::generateKeyPair");

    #ifdef DEBUG_PRINTS_LEVEL_2
    std::cout << "[DSA_generateKeyPair][debug]Generating DSA key pair..." << std::endl;
    #endif
    int L = key_length; // 密钥长度
    int N; // 签名的位数
    {
        ZoneScopedN("DSA::generateKeyPair::Key Length Selection");
        std::map<int, int>N_map = { {1024,160},{2048,224},{3072,256} };
        if (N_map.find(L) != N_map.end()) {
            N = N_map[L];
        }
        else {
            throw std::runtime_error("Invalid key length");
        }
    }


    // 验证公共参数有效性
    {
        ZoneScopedN("DSA::generateKeyPair::Parameter Validation");
        if (!publicParameters.is_valid) {
            DSA::generatepublicparameters(publicParameters); // 内部已插桩
        }
    }

    // 初始化临时变量
    mpz_t x, y;
    {
        ZoneScopedN("DSA::generateKeyPair::GMP Variable Initialization");
        mpz_inits(x, y, nullptr);
    }

    // 初始化随机数生成器
    gmp_randstate_t state;
    {
        ZoneScopedN("DSA::generateKeyPair::Random State Initialization");
        gmp_randinit_default(state);
        gmp_randseed_ui(state, static_cast<unsigned long>(time(nullptr)));
    }

    // ---------------------------------
    // 1. 选择私钥x。x为一个随机生成的整数
    // ---------------------------------
    {
        ZoneScopedN("DSA::generateKeyPair::Choose Private Key X");
        mpz_urandomb(x, state, N);
    }

    // --------------------------
    // 2. 计算公钥y: y = g^x mod p
    // --------------------------
    {
        ZoneScopedN("DSA::generateKeyPair::Calculate Public Key Y");
        mpz_powm(y, publicParameters.g, x, publicParameters.p);
    }

    // 保存密钥对
    {
        ZoneScopedN("DSA::generateKeyPair::Store Key Pair and Cleanup");
        publicKey = DSA_publickey(y);
        privateKey = DSA_privatekey(x);

        mpz_clears(x, y, nullptr);
        gmp_randclear(state);
    }
}

template<int key_length>
DSA_signature DSA<key_length>::sign(const PlainMessage& message,
    const DSA_privatekey& privateKey,
    const DSA_publicparameter& publicParameters) {
    // 为 sign 函数添加 Tracy 区域
    ZoneScopedN("DSA::sign");

    #ifdef DEBUG_PRINTS_LEVEL_2
    std::cout << "[DSA_sign][debug] Starting DSA signature generation..." << std::endl;
    #endif

    // 验证公共参数有效性
    {
        ZoneScopedN("DSA::sign::Parameter Validation");
        if (!publicParameters.is_valid) {
            throw std::invalid_argument("Invalid DSA public parameters");
        }
    }

    // 初始化GMP变量
    mpz_t hash_m, k, r, s, kinv, temp;
    gmp_randstate_t state;
    {
        ZoneScopedN("DSA::sign::GMP Variable and Random State Initialization");
        mpz_inits(hash_m, k, r, s, kinv, temp, nullptr);
        gmp_randinit_default(state);
    }

    try {
        // 设置随机种子
        {
            ZoneScopedN("DSA::sign::Random Seed Setting");
            gmp_randseed_ui(state, static_cast<unsigned long>(time(nullptr)));
        }

        // -----------------------------
        // 1. 将消息哈希并转换为整数hash_m
        // -----------------------------
        {
            ZoneScopedN("DSA::sign::Hash Message to Integer");
            std::string hex_str;
            hash_value_to_hex_string(message.get_hash_data(), hex_str);
            int ret = mpz_set_str(hash_m, hex_str.c_str(), 16);
            if (ret != 0) {
                throw std::runtime_error("Invalid hash format");
            }
        }

        #ifdef DEBUG_PRINTS_LEVEL_3
        gmp_printf("[DSA_sign] Hashed message: %Zx\n", hash_m);
        #endif

        // 步骤2-3：生成临时密钥k并计算r
        int attempts = 0;
        constexpr int MAX_ATTEMPTS = 100;
        {
            ZoneScopedN("DSA::sign::Generate K and Calculate R Loop");
            do {
                // -----------------------------------
                // 2. 产生一个随机数k, 满足k ∈ [1, q-1]
                // -----------------------------------
                {
                    ZoneScopedN("DSA::sign::Generate K");
                    mpz_urandomm(k, state, publicParameters.q);
                }
                #ifdef DEBUG_PRINTS_LEVEL_3
                gmp_printf("[DSA_sign] Generated temporary key k: %Zx\n", k);
                #endif
                // ----------------------------
                // 3. 计算r = (g^k mod p) mod q
                // ----------------------------
                {
                    ZoneScopedN("DSA::sign::Calculate R");
                    mpz_powm(r, publicParameters.g, k, publicParameters.p);
                    mpz_mod(r, r, publicParameters.q);
                }

                if (++attempts > MAX_ATTEMPTS) {
                    throw std::runtime_error("Failed to generate valid k after multiple attempts");
                }
            } while (mpz_sgn(r) == 0);   // 确保r≠0
        }

        #ifdef DEBUG_PRINTS_LEVEL_3
        gmp_printf("[DSA_sign] Generated r: %Zx\n", r);
        #endif

        // 计算k的模逆： kinv = k^-1 mod q
        {
            ZoneScopedN("DSA::sign::Calculate K Inverse");
            if (mpz_invert(kinv, k, publicParameters.q) == 0) {
                throw std::runtime_error("Failed to compute k inverse");
            }
        }

        // --------------------------------------
        // 4. 计算s = (kinv * (hash + x*r)) mod q
        // --------------------------------------
        {
            ZoneScopedN("DSA::sign::Calculate S");
            mpz_mul(temp, privateKey.x, r);              // x*r
            mpz_add(temp, temp, hash_m);                 // hash + x*r
            mpz_mul(s, temp, kinv);                      // kinv*(hash + x*r)
            mpz_mod(s, s, publicParameters.q);           // mod q
        }
        #ifdef DEBUG_PRINTS_LEVEL_3
        gmp_printf("[DSA_sign] hash_m: %Zd\n", hash_m);
        gmp_printf("[DSA_sign] kinv: %Zd\n", kinv);
        gmp_printf("[DSA_sign] temp: %Zd\n", temp);
        gmp_printf("[DSA_sign] s: %Zd\n", s);
        #endif
        // 验证s≠0
        {
            ZoneScopedN("DSA::sign::Verify S Non-Zero");
            if (mpz_sgn(s) == 0) {
                throw std::runtime_error("Generated invalid s=0");
            }
        }

        #ifdef DEBUG_PRINTS_LEVEL_3
        gmp_printf("[DSA_sign] Generated s: %Zx\n", s);
        #endif

        // 构造签名对象
        DSA_signature signature{ r, s };

        // 清理资源
        {
            ZoneScopedN("DSA::sign::Cleanup Resources");
            mpz_clears(hash_m, k, r, s, kinv, temp, nullptr);
            gmp_randclear(state);
        }

        #ifdef DEBUG_PRINTS_LEVEL_2
        std::cout << "[DSA_sign][debug] Signature generated successfully" << std::endl;
        #endif
        return signature;

    }
    catch (...) {
        // 异常时清理资源
        ZoneScopedN("DSA::sign::Exception Cleanup");
        mpz_clears(hash_m, k, r, s, kinv, temp, nullptr);
        gmp_randclear(state);
        throw;
    }
}

template<int key_length>
bool DSA<key_length>::verify(const PlainMessage& message, const DSA_publickey& publicKey, const DSA_publicparameter& publicParameters, const DSA_signature& signature) {
    // 为 verify 函数添加 Tracy 区域
    ZoneScopedN("DSA::verify");

    // 验证公共参数有效性
    {
        ZoneScopedN("DSA::verify::Parameter Validation");
        if (!publicParameters.is_valid) {
            throw std::invalid_argument("Invalid public parameters");
        }
    }

    // 初始化GMP变量
    mpz_t hash_m, w, u1, u2, v, temp;
    {
        ZoneScopedN("DSA::verify::GMP Variable Initialization");
        mpz_inits(hash_m, w, u1, u2, v, temp, nullptr);
    }

    try {
        // 验证签名范围 (0 < r,s < q)
        {
            ZoneScopedN("DSA::verify::Signature Range Check");
            if (
                mpz_cmp(signature.r, publicParameters.q) >= 0 ||
                mpz_cmp(signature.s, publicParameters.q) >= 0
                ) {
                #ifdef DEBUG_PRINTS_LEVEL_2
                std::cerr << "[DSA_verify][debug] Invalid signature range" << std::endl;
                #endif
                return false;
            }
        }

        // -----------------------------
        // 1. 将消息哈希并转换为整数hash_m
        // -----------------------------
        {
            ZoneScopedN("DSA::verify::Hash Message to Integer");
            std::string hash_str;
            hash_value_to_hex_string(message.get_hash_data(), hash_str);
            int ret = mpz_set_str(hash_m, hash_str.c_str(), 16);
            if (ret != 0) {
                throw std::runtime_error("Invalid hash format");
            }
        }
        #ifdef DEBUG_PRINTS_LEVEL_2
        gmp_printf("[DSA_verify][debug] Hash value: %Zd\n", hash_m);
        #endif
        // -----------------------
        // 2. 计算w = s^{-1} mod q
        // -----------------------
        {
            ZoneScopedN("DSA::verify::Calculate W (s_inverse)");
            if (mpz_invert(w, signature.s, publicParameters.q) == 0) {
                #ifdef DEBUG_PRINTS_LEVEL_2
                gmp_printf("[DSA_verify][debug] Cannot invert s: %Zd\n", signature.s);
                gmp_printf("[DSA_verify][debug] Cannot invert r: %Zd\n", signature.r);
                #endif
                return false;
            }
        }

        // ------------------------------
        // 3. 计算u1 = (hash_m * w) mod q
        // ------------------------------
        {
            ZoneScopedN("DSA::verify::Calculate U1");
            mpz_mul(u1, hash_m, w);
            mpz_mod(u1, u1, publicParameters.q);
        }

        // -------------------------
        // 4. 计算u2 = (r * w) mod q
        // -------------------------
        {
            ZoneScopedN("DSA::verify::Calculate U2");
            mpz_mul(u2, signature.r, w);
            mpz_mod(u2, u2, publicParameters.q);
        }

        // ------------------------------------
        // 5. 计算v = (g^u1 * y^u2 mod p) mod q
        // ------------------------------------
        {
            ZoneScopedN("DSA::verify::Calculate V");
            mpz_powm(temp, publicParameters.g, u1, publicParameters.p);    // g^u1 mod p
            // 这里u2被复用为结果，先保存原始u2的值
            mpz_t u2_orig;
            mpz_init_set(u2_orig, u2); // 保存u2的原始值
            mpz_powm(u2, publicKey.y, u2_orig, publicParameters.p);        // y^u2 mod p
            mpz_clear(u2_orig); // 清理临时变量

            mpz_mul(v, temp, u2);                                 // g^u1 * y^u2
            mpz_mod(v, v, publicParameters.p);                   // mod p
            mpz_mod(v, v, publicParameters.q);                   // mod q
        }
        #ifdef DEBUG_PRINTS_LEVEL_3
        gmp_printf("[DSA_verify][debug] Intermediate values:\n");

        gmp_printf("hash_m: %Zd\n", hash_m);
        gmp_printf("u1: %Zd\nu2: %Zd\nv: %Zd\n", u1, u2, v);
        gmp_printf("w: %Zd\n", w);
        #endif
        // 验证v == r
        bool isValid;
        {
            ZoneScopedN("DSA::verify::Compare V and R");
            isValid = (mpz_cmp(v, signature.r) == 0);
        }

        #ifdef DEBUG_PRINTS_LEVEL_3
        gmp_printf("[DSA_verify][debug] Verification result: %d\n", isValid);
        gmp_printf("Computed v: %Zd\nOriginal r: %Zd\n", v, signature.r);
        #endif

        // 清理资源
        {
            ZoneScopedN("DSA::verify::Cleanup Resources");
            mpz_clears(hash_m, w, u1, u2, v, temp, nullptr);
            return isValid;
        }

    }
    catch (...) {
        // 异常时确保清理资源
        ZoneScopedN("DSA::verify::Exception Cleanup");
        mpz_clears(hash_m, w, u1, u2, v, temp, nullptr);
        throw;
    }
}

template class DSA<1024>;
template class DSA<2048>;
template class DSA<3072>;